import numpy as np
import acl

ACL_SUCCESS = 0
ACL_MEMCPY_HOST_TO_HOST = 0
ACL_MEMCPY_HOST_TO_DEVICE = 1
ACL_MEMCPY_DEVICE_TO_HOST = 2
ACL_MEM_MALLOC_HUGE_FIRST = 0
ACL_DEVICE, ACL_HOST = 0, 1
NPY_FLOAT16 = 23


def main():
    device_id = 0
    # acl资源初始化
    ret = acl.init()
    if ret != ACL_SUCCESS:
        print("acl.init failed ret = {}".format(ret))
        exit(1)
    ret = acl.rt.set_device(device_id)
    if ret != ACL_SUCCESS:
        print("acl.rt.set_device failed ret = {}".format(ret))
        exit(1)
    context, ret = acl.rt.create_context(device_id)
    if ret != ACL_SUCCESS:
        print("acl.rt.create_context failed ret = {}".format(ret))
        exit(1)


    # 准备CPU运行内存数据
    random_matrix = np.random.rand(3, 1920, 1080).astype(np.float16)
    numpyDataByte = random_matrix.tobytes()
    numpyDataPtr = acl.util.bytes_to_ptr(numpyDataByte)

    # 准备NPU上显存空间
    deviceDataPtr, ret = acl.rt.malloc(3*1920*1080*2, ACL_MEM_MALLOC_HUGE_FIRST)

    # CPU内存拷贝到NPU上
    runMode = acl.rt.get_run_mode()         # 判断当前程序是运行在cpu上还是npu上
    if runMode == ACL_DEVICE:
        kind = ACL_MEMCPY_DEVICE_TO_DEVICE  # 如果程序是NPU上，就是把NPU显存上的内存拷贝到NPU显存上
    else:
        kind = ACL_MEMCPY_HOST_TO_DEVICE    # 如果程序是CPU上，就是运行内存中的内容拷贝到NPU显存上
    ret = acl.rt.memcpy(deviceDataPtr, 3*1920*1080*2, numpyDataPtr, 3*1920*1080*2, kind) # CPU内存拷贝到NPU上


    # 准备CPU运行内存上空间
    outDataPtr, ret = acl.rt.malloc_host(3*1920*1080*2)

    # NPU内存上的内容拷贝到CPU上
    if runMode == ACL_DEVICE:
        kind = ACL_MEMCPY_DEVICE_TO_DEVICE  # 如果程序是NPU上，就是把NPU显存上的内存拷贝到NPU显存上
    else:
        kind = ACL_MEMCPY_DEVICE_TO_HOST    # 如果程序是CPU上，就是NPU显存中的内容拷贝到运行内存上
    ret = acl.rt.memcpy(outDataPtr, 3*1920*1080*2, deviceDataPtr, 3*1920*1080*2, kind) # NPU内存拷贝到CPU上


    # 将输出数据转成numpy
    outByteData = acl.util.ptr_to_bytes(outDataPtr, 3 * 1920 * 1080 * 2)
    outNumpyData = np.frombuffer(outByteData, dtype=np.float16).reshape(3, 1920, 1080)

    # 判断outNumpyData和random_matrix是不是一样的
    print("outNumpyData==random_matrix is", np.array_equal(outNumpyData, random_matrix))


    # 释放acl资源
    ret= acl.rt.free(deviceDataPtr)
    if ret != ACL_SUCCESS:
        print("acl.rt.free ret = {}".format(ret))
    ret = acl.rt.destroy_context(context)
    if ret != ACL_SUCCESS:
        print("acl.rt.destroy_context ret = {}".format(ret))
    ret = acl.rt.reset_device(device_id)
    if ret != ACL_SUCCESS:
        print("acl.rt.reset_device ret = {}".format(ret))
    ret = acl.finalize()
    if ret != ACL_SUCCESS:
        print("acl.finalize failed ret = {}".format(ret))
    print("run success!")

if __name__ == '__main__':
    main()
